import numpy as np
import pandas as pd
from coommon import *

class TwoLayerNet:
    def __init__(self,input_size, hidden_size, output_size, weight_init_std=0.01):
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

    def predict(self, x):
        W1, W2 = self.params['W1'], self.params['W2']
        b1, b2 = self.params['b1'], self.params['b2']

        a1 = np.dot(x,W1) + b1
        z1 = sigmoid(a1)

        a2 = np.dot(z1,W2) + b2
        y = softmax(a2)

        return y

    def accuracy(self,x,t):
        y = self.predict(x)
        y = np.argmax(y,axis=1)
        t = t.reshape(-1)

        accuracy = np.sum(y==t)/t.shape[0] * 1.0
        return accuracy

    def loss(self,x,t):
        y = self.predict(x)
        return cross_entropy_error(y,t)

    def numerical_gradient(self,x,t):
        loss_w = lambda w:self.loss(x,t)

        grads = { }
        grads['W1'] = numerical_gradient(loss_w,self.params['W1'])
        grads['b1'] = numerical_gradient(loss_w, self.params['b1'])
        grads['W2'] = numerical_gradient(loss_w, self.params['W2'])
        grads['b2'] = numerical_gradient(loss_w,self.params['b2'])

        return grads

if __name__ == '__main__':
    #1加载数据
    X_train, X_test, y_train, y_test = get_data()
    #2创建网络
    net = TwoLayerNet(input_size=784,hidden_size=100,output_size=10)
    #3设置超参数
    lr = 0.05
    batch_size = 100
    num_epochs = 20
    n = X_train.shape[0]
    iter_per_epoch = int(np.ceil(n / batch_size))
    iters_num = int(iter_per_epoch * num_epochs)
    train_loss_list = []
    train_acc_list = []
    test_acc_list= []
    #4循环迭代
    for i in range(iters_num):
        # 4.1 随机选取批量数据
        batch_idx = np.random.choice(n,batch_size)
        x_batch = X_train[batch_idx]
        t_batch = y_train[batch_idx]
        #4.2 计算梯度
        grads = net.numerical_gradient(x_batch,t_batch)
        # 更新参数
        for key in  ('W1','b1','W2','b2'):
            net.params[key] -= lr * grads[key]
        # 计算并保存损失
        train_loss_list.append(net.loss(x_batch,t_batch))
        # 按照轮次（epoch）输出准确率
        if i % iter_per_epoch == 0:
            train_acc_list.append(net.accuracy(X_train,y_train))
            test_acc_list.append(net.accuracy(X_test, y_test))
            print('epoch:{},train_acc:{},test_acc:{}'.format(i//iter_per_epoch,train_acc_list[-1],test_acc_list[-1]))
    #5画图
    x = np.array(len(test_acc_list))
    plt.plot(x,train_acc_list,label='train_acc')
    plt.plot(x,test_acc_list,label='test_acc',line_style='dashed')
    plt.legend()
    plt.show()


